
from service.wrap import stream_wrapper
from fastapi import FastAPI, Body
from fastapi.responses import StreamingResponse
import uvicorn
import time


@stream_wrapper
async def do_works(
        name: str = Body('hello'),
        number: int = 10
):
    for n in range(number):
        time.sleep(n)
        yield {'name': name, 'n': n}

app = FastAPI()

@app.post("/work")
async def work_api(
    name: str = Body(default='hello'),
    number: int = Body(default=10)
):
    return StreamingResponse(do_works(name, number), media_type="text/event-stream")


if __name__ == "__main__":
    uvicorn.run(host="0.0.0.0", port=8000, app=app)
